# src/data/market_data.py
import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Any
from datetime import datetime, timedelta
import yfinance as yf
from sklearn.preprocessing import StandardScaler
import warnings

class MarketDataProcessor:
    """市场数据处理器"""
    
    def __init__(self):
        self.scaler = StandardScaler()
        self.price_data = {}
        self.returns_data = {}
        
    def fetch_stock_data(self, symbols: List[str], 
                        start_date: str, end_date: str) -> pd.DataFrame:
        """获取股票数据"""
        try:
            data = yf.download(symbols, start=start_date, end=end_date)
            return data['Adj Close']
        except Exception as e:
            warnings.warn(f"获取数据失败: {e}")
            return pd.DataFrame()
    
    def calculate_returns(self, price_data: pd.DataFrame, 
                         method: str = 'log') -> pd.DataFrame:
        """计算收益率"""
        if method == 'log':
            returns = np.log(price_data / price_data.shift(1))
        elif method == 'simple':
            returns = price_data.pct_change()
        else:
            raise ValueError(f"不支持的收益率计算方法: {method}")
        
        return returns.dropna()
    
    def calculate_ewma_covariance(self, returns: pd.DataFrame, 
                                 lambda_param: float = 0.94) -> np.ndarray:
        """计算EWMA协方差矩阵"""
        n, m = returns.shape
        cov_matrix = np.zeros((m, m))
        
        # 初始化协方差矩阵
        cov_matrix = np.cov(returns.iloc[:30].T)
        
        # 递归更新EWMA协方差
        for t in range(30, n):
            r_t = returns.iloc[t].values.reshape(-1, 1)
            cov_matrix = lambda_param * cov_matrix + (1 - lambda_param) * np.dot(r_t, r_t.T)
        
        return cov_matrix
    
    def garch_volatility_forecast(self, returns: pd.Series, 
                                 forecast_horizon: int = 1) -> float:
        """GARCH波动率预测（简化版）"""
        try:
            from arch import arch_model
            
            # 拟合GARCH模型
            model = arch_model(returns * 100, vol='GARCH', p=1, q=1)
            fitted_model = model.fit(disp='off')
            
            # 预测波动率
            forecast = fitted_model.forecast(horizon=forecast_horizon)
            predicted_vol = np.sqrt(forecast.variance.iloc[-1, 0]) / 100
            
            return predicted_vol
            
        except ImportError:
            warnings.warn("arch包未安装，使用滚动标准差替代")
            return returns.rolling(window=30).std().iloc[-1]
    
    def calculate_technical_indicators(self, price_data: pd.DataFrame) -> pd.DataFrame:
        """计算技术指标"""
        indicators = pd.DataFrame(index=price_data.index)
        
        for symbol in price_data.columns:
            prices = price_data[symbol]
            
            # RSI
            delta = prices.diff()
            gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
            loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
            rs = gain / loss
            indicators[f'{symbol}_RSI'] = 100 - (100 / (1 + rs))
            
            # MACD
            exp1 = prices.ewm(span=12).mean()
            exp2 = prices.ewm(span=26).mean()
            macd = exp1 - exp2
            signal = macd.ewm(span=9).mean()
            indicators[f'{symbol}_MACD'] = macd - signal
            
            # 布林带
            sma20 = prices.rolling(window=20).mean()
            std20 = prices.rolling(window=20).std()
            indicators[f'{symbol}_BB_upper'] = sma20 + (std20 * 2)
            indicators[f'{symbol}_BB_lower'] = sma20 - (std20 * 2)
            indicators[f'{symbol}_BB_position'] = (prices - sma20) / (std20 * 2)
        
        return indicators.dropna()
    
    def get_esg_scores(self, symbols: List[str]) -> Dict[str, float]:
        """获取ESG评分（模拟数据）"""
        # 在实际应用中，这里应该连接到ESG数据提供商API
        esg_scores = {}
        
        # 模拟ESG评分
        np.random.seed(42)  # 保证可重复性
        for symbol in symbols:
            # 基于symbol生成相对稳定的ESG评分
            score = 50 + 30 * np.sin(hash(symbol) % 100) + np.random.normal(0, 5)
            esg_scores[symbol] = max(0, min(100, score))
        
        return esg_scores
    
    def prepare_portfolio_data(self, symbols: List[str], 
                              start_date: str, end_date: str) -> Dict[str, Any]:
        """准备投资组合数据"""
        # 获取价格数据
        price_data = self.fetch_stock_data(symbols, start_date, end_date)
        
        if price_data.empty:
            raise ValueError("无法获取价格数据")
        
        # 计算收益率
        returns = self.calculate_returns(price_data)
        
        # 计算协方差矩阵
        cov_matrix = self.calculate_ewma_covariance(returns)
        
        # 计算期望收益
        expected_returns = returns.mean().values
        
        # 获取ESG评分
        esg_scores = self.get_esg_scores(symbols)
        esg_array = np.array([esg_scores[symbol] for symbol in symbols])
        
        # 计算技术指标
        technical_indicators = self.calculate_technical_indicators(price_data)
        
        return {
            'price_data': price_data,
            'returns': returns,
            'expected_returns': expected_returns,
            'cov_matrix': cov_matrix,
            'esg_scores': esg_array,
            'technical_indicators': technical_indicators,
            'symbols': symbols
        }
    
    def get_market_indicators(self, date: datetime) -> Dict[str, float]:
        """获取市场指标"""
        # 模拟市场指标数据
        # 在实际应用中，应该连接到实时数据源
        
        return {
            'vix': np.random.normal(20, 5),  # 波动率指数
            'market_return': np.random.normal(0.001, 0.02),  # 市场收益
            'market_volatility': np.random.normal(0.15, 0.03),  # 市场波动率
            'rsi': np.random.normal(50, 15),  # 市场RSI
            'macd': np.random.normal(0, 0.5),  # 市场MACD
            'volume_ratio': np.random.normal(1.0, 0.2)  # 成交量比率
        }
